import argparse, os, time
from pprint import pprint
import pandas as pd


parser = argparse.ArgumentParser()
parser.add_argument('record_path', type=str,
                    help='the path to load pretrained models and save exp results')
parser.add_argument('--sample_path', type=str, default='',
                    help='the path to save exp results')
parser.add_argument('--dataset', type=str, default='coco',
                    help='training dataset')
parser.add_argument('--img_size', type=int, default=128, help='image size')
parser.add_argument("--gpu", default="0", type=str,
        help='GPU to use (leave blank for CPU only)')
parser.add_argument("--end", default=0, type=int,
        help='The maximum epoch to use')
parser.add_argument("--begin", default=0, type=int,
        help='The minimum epoch to use')
# parser.add_argument("--interval", default=5, type=int,
#         help='The interval to use')
parser.add_argument('-G', '--generate', action="store_true", default=False,
        help='to generate images')
parser.add_argument('--seed', type=int, default=0, help='random seed')
parser.add_argument('--test', action="store_true", default=False, help='random seed')
args = parser.parse_args()

ISresult = dict()
FIDresult = dict()
FIDtest_result = dict()
DS_result = dict()
SceneFID_result = dict()
PRC_result = dict()
PRCOBJ_result = dict()
mask_entropy_result = dict()
raw_mask_entropy_result = dict()

exp_sample_path = os.path.join([args.record_path, args.sample_path][len(args.sample_path)>0], 'samples')
for epoch in list(range(args.begin, args.end+1))+['outside']:
    if not os.path.exists( os.path.join(args.record_path, "model", str(epoch)) ):
        continue
    else:
        print(f"I come to iter {epoch}.")
    model_path = os.path.join(args.record_path, "model", str(epoch))
    epoch_path = os.path.join(exp_sample_path, f"iter_{epoch}")

    print(f"Using model in {model_path}")
    print(f"Will save results in {epoch_path}")
    
    repeat = 5
    save_path = os.path.join(epoch_path, f"{args.dataset}_{args.img_size}_repeat{repeat}")
    img_save_path = os.path.join(save_path, "img" )
    crop_save_path = os.path.join(save_path, "cropped_224" )

    # generate_cmd = f"~/anaconda3/envs/zl8/bin/python test.py --dataset {args.dataset} --model_path {model_path} --sample_path {epoch_path} -r {repeat} --gpu {args.gpu} --img_size {args.img_size} --seed {args.seed}"
    generate_cmd = f"python test.py --dataset {args.dataset} --model_path {model_path} --sample_path {epoch_path} -r {repeat} --gpu {args.gpu} --img_size {args.img_size} --seed {args.seed}"
    print(generate_cmd)
    if args.generate:
        if not os.path.exists(epoch_path):
            print(generate_cmd)
            value = os.popen(generate_cmd).readlines()[-1]
            print(value)
            exec('tmp = {'+value+'}' )
            print(tmp)
            DS_result[epoch] = tmp['Diversity_Score']
            mask_entropy_result[epoch] = tmp["mask_entropy"]
            raw_mask_entropy_result[epoch] = tmp["raw_mask_entropy"]
        else: 
            print(f"{epoch_path} exists.")
            time.sleep(3)
            DS_result[epoch] = -1
            mask_entropy_result[epoch] = -1
            raw_mask_entropy_result[epoch] = -1


    # FID
    try:
        val_set_path = f"datasets/{args.dataset}/val_{args.img_size}/"
        # sample_path = os.path.join(epoch_path, f"{args.dataset}{args.img_size}_repeat5_thres2.0/") 
        fid_cmd = f"~/.conda/envs/tf24/bin/python scores/FID.py {val_set_path} {img_save_path} --gpu {args.gpu} --lowprofile"
        print(fid_cmd)
        if args.test:
            pass
        else:
            value = os.popen(fid_cmd).readlines()[-1].split(' ')[-1]    
            FIDresult[epoch] = float(value)
            print(value)
    except:
        FIDresult[epoch] = -1

    # IS
    try:
        is_cmd = f"~/.conda/envs/tf24/bin/python scores/InceptionScore.py {img_save_path} --gpu {args.gpu}"
        print(is_cmd)
        if args.test:
            pass
        else:
            value = os.popen(is_cmd).readlines()[-1].split(' ')[-1]
            ISresult[epoch] = value
            print(value)
    except:
        ISresult[epoch] = -1

    # # FID test
    # try:
    #     test_set_path = f"datasets/{args.dataset}/test_{args.img_size}/"
    #     fid_cmd = f"~/.conda/envs/tf24/bin/python scores/FID.py {test_set_path} {sample_path} --gpu {args.gpu}"
    #     print(fid_cmd)
    #     value = os.popen(fid_cmd).readlines()[-1].split(' ')[-1]
    #     FIDtest_result[epoch] = float(value[:8])
    #     print(value)
    # except:
    #     FIDtest_result[epoch] = -1

    # SceneFID
    try:
        # SceneFID_cmd = f"~/anaconda3/envs/zl6/bin/python scores/SceneFID.py --model_path {model_path} --intermediate_path {epoch_path} --dataset {args.dataset} --img_size {args.img_size} --gpu {args.gpu}"
        val_crop_path = f"datasets/{args.dataset}/val_{args.img_size}_cropped_224"
        SceneFID_cmd = f"~/.conda/envs/tf24/bin/python scores/FID.py {val_crop_path} {crop_save_path} --gpu {args.gpu} --lowprofile"

        print(SceneFID_cmd)
        if args.test:
            pass
        else:
            value = os.popen(SceneFID_cmd).readlines()[-1].split(' ')[-1]
            SceneFID_result[epoch] = float(value[:8])
            print(value)
    except:
        SceneFID_result[epoch] = -1


    print("I come to repeat 1 generation")
    # repeat 1 for prc and prc obj
    repeat1 = 1
    save_path_repeat1 = os.path.join(epoch_path, f"{args.dataset}_{args.img_size}_repeat{repeat1}")
    img_save_path_repeat1 = os.path.join(save_path_repeat1, "img" )
    crop_save_path_repeat1 = os.path.join(save_path_repeat1, "cropped_224" )

    # generate_cmd = f"~/anaconda3/envs/zl8/bin/python test.py --dataset {args.dataset} --model_path {model_path} --sample_path {epoch_path} -r {repeat} --gpu {args.gpu} --img_size {args.img_size} --seed {args.seed}"
    generate_cmd = f"python test.py --dataset {args.dataset} --model_path {model_path} --sample_path {epoch_path} -r {repeat1} --gpu {args.gpu} --img_size {args.img_size} --seed {args.seed}"
    print(generate_cmd)
    if args.generate:
        if not os.path.exists(save_path_repeat1):
            print(generate_cmd)
            os.popen(generate_cmd).readlines()
        else: 
            print(f"{save_path_repeat1} exists.")
            time.sleep(3)

    # prc
    try:
        layout_path = f"{args.dataset}.txt"
        val_set_path = f"datasets/{args.dataset}/val_{args.img_size}/"
        # num_samples = {'coco':3097, 'vg':5064}[args.dataset.lower()]
        num_samples = int(os.popen(f"ls -lR {val_set_path} | grep '^-' | wc -l").readlines()[-1])

        PRC_cmd = f"python scores/PRC.py --path_real {val_set_path} --layout_real {layout_path} --path_fake {img_save_path_repeat1} --layout_fake {layout_path} --num_samples {num_samples} --gpu {args.gpu}"

        print(PRC_cmd)
        # value = os.popen(PRC_cmd).readlines()[-1]
        # PRC_result[epoch] = value
        # print("PRC   ", value)
    except:
        PRC_result[epoch] = -1

    # #prc obj
    # try:
    #     val_crop_path = f"datasets/{args.dataset}/val_{args.img_size}_cropped_224"
    #     num_samples = int(os.popen(f"ls -lR {val_set_path} | grep '^-' | wc -l").readlines()[-1])
    #     PRCOBJ_cmd = f"python scores/PRC_object.py --path_real {val_crop_path} --path_fake {crop_save_path_repeat1} --num_samples {num_samples} --gpu {args.gpu}"
    #     print(PRCOBJ_cmd)
    #     if args.test:
    #         pass
    #     else:
    #         value = os.popen(PRCOBJ_cmd).readlines()[-1]
    #         PRCOBJ_result[epoch] = value
    #         print("PRC obj   ", value)
    # except:
    #     PRCOBJ_result[epoch] = -1

    
    # tmp_tab = pd.DataFrame.from_dict({"IS":ISresult, "FID":FIDresult, "FID_test":FIDtest_result, "DS":DS_result, "SceneFID":SceneFID_result})
    tmp_tab = pd.DataFrame.from_dict({
     "FID":FIDresult,
     "SceneFID":SceneFID_result,
        "IS":ISresult,
     "DS": DS_result,
     "mask_entropy": mask_entropy_result,
     "raw_mask_entropy": raw_mask_entropy_result,
     })
    print("tmp_tab", tmp_tab)
    file_path = os.path.join(exp_sample_path, 'pd.csv')
    if os.path.exists(file_path):
        tab_dict = pd.read_csv(file_path, index_col=0).to_dict(orient='index')
        # tab_dict.update(tmp_tab.to_dict(orient='index'))
        print("tab_dict", tab_dict)
        for k, v in tmp_tab.to_dict(orient='index').items():
            tab_dict[k] = v
        tab = pd.DataFrame.from_dict(tab_dict, orient='index')
    else:
        tab = tmp_tab
    tab.to_csv(file_path)
    
    print(f"epoch_{epoch}")
    print(tab)
    print(tab.to_markdown())

# python scores/Evaluation.py experiments/coco_128_20w_double_cg_DC/2021-12-10_21\:59\:55/  --gpu 4 -G --img_size 128 --dataset coco --end 200000 --begin 0
# python scores/Evaluation.py experiments/coco_128_20w_replicationpad_fixed_scale_noise/2021-11-13_10\:55\:35/  --gpu 0 -G --img_size 128 --dataset coco --end 30000 --begin 0
# python scores/Evaluation.py experiments/coco_128_20w_replicationpad/2021-11-11_09\:13\:44/  --gpu 4 -G --img_size 128 --dataset coco --end 40000 --begin 0
# python scores/Evaluation.py experiments/coco_128_reproduce_lama_20k_coloraug_no_warmup/2021-11-07_17\:48\:31/  --gpu 4 -G --img_size 128 --dataset coco --end 70001
# python scores/Evaluation.py experiments/coco_128_reproduce_lama_20k_coloraug/2021-11-04_21\:19\:20/ --gpu 4 -G --img_size 128 --dataset coco --end 200000
# python scores/Evaluation.py experiments/coco_128_reproduce_lama_20k_coloraug/2021-11-04_21\:19\:20/ --gpu 4 -G --img_size 128 --dataset coco
# python scores/Evaluation.py experiments/coco_64_reproduce_lama_20k/2021-10-29_21:00:50/ --gpu 4 -G --img_size 64 --dataset coco
# python scores/Evaluation.py experiments/coco_64_reproduce_lama/2021-10-24_20:13:30/ --gpu 7 -G --img_size 64 --dataset coco
# python scores/Evaluation.py experiments/coco_64_1018/2021-10-18_20\:05\:05/ --gpu 0 -G --img_size 64 --dataset coco
# experiments/coco_64_1004_continue/2021-10-05_16:19:52/samples/iter_outside
# python scores/Evaluation.py experiments/coco_64_1004_continue/2021-10-05_16\:19\:52/ --gpu 7 -G --img_size 64
# python scores/Evaluation.py experiments/coco_64_1001_continue/2021-10-01_21\:26\:01/ --gpu 6 -G --img_size 64

# python scores/Evaluation.py experiments/coco0616/ --gpu 6 -G --img_size 64 --epoch 190

# python scores/Evaluation.py experiments/coco0805/ --gpu 0 -G --img_size 128 --epoch 180

# python scores/Evaluation.py experiments/vg0829/ --dataset vg -G --img_size 64 --begin_epoch 145


# python scores/Evaluation.py experiments/vg0829/ --dataset vg -G --img_size 64 --begin_epoch 145

# python scores/Evaluation.py experiments/coco0916/ --dataset coco -G --img_size 256 --begin_epoch 130 --gpu 2

# python scores/Evaluation.py experiments/coco0916/ --dataset coco -G --img_size 256 --begin_epoch 120 --end_epoch 120 --gpu 2

# python scores/Evaluation.py experiments/vg0927/ --dataset vg -G --img_size 256 --begin_epoch 200 --gpu 6

# python scores/Evaluation.py experiments/coco0727/ --dataset coco -G --img_size 64 --begin_epoch 200 --gpu 2

# python scores/Evaluation.py experiments/coco1006_128_LostGANv1_with_adjust/ --dataset coco -G --img_size 128 --begin_epoch 175 --end_epoch 175 --gpu 1

# python scores/Evaluation.py experiments/coco1118_only_pixel_query/ --dataset coco -G --img_size 128 --begin_epoch 200 --gpu 1

# python scores/Evaluation.py rebuttal_two_stage_2/ --dataset coco -G --img_size 128 --begin_epoch 5 --end_epoch 5 --gpu 7

# python scores/Evaluation.py rebuttal_3/ --dataset coco -G --img_size 128 --begin_epoch 12 --end_epoch 12 --gpu 7
# python scores/Evaluation.py rebuttal_3/ --dataset coco -G --img_size 128 --begin_epoch 11 --end_epoch 11 --gpu 4

# python scores/Evaluation.py rebuttal_3/ --dataset coco -G --img_size 128 --begin_epoch 11 --end_epoch 11 --gpu 4

# python scores/Evaluation.py experiments/R3W3_06216_1117/ --dataset coco -G --img_size 128 --begin_epoch 10 --end_epoch 200 --interval 10  --gpu 7

